Checkpoint the curriculum data sampler's own RNG state - #8406
ebarkhordar wants to merge 1 commit into
Conversation
DeepSpeedDataSampler draws from self.np_rng, but state_dict() saved np.random.get_state() and load_state_dict() restored it. That is the global legacy RandomState, so the sampler's generator was never in the checkpoint and a resumed run replayed its sampling stream from the seed. Save and restore self.np_rng.bit_generator.state instead. A state dict written by an older version carries the old tuple, which still goes through np.random.set_state. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8384ae8662
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| resumed = _curriculum_data_sampler(tmp_path) | ||
| resumed.load_state_dict(saved.state_dict()) | ||
| next_draws = [resumed.sample_from_clusters().tolist() for _ in range(3)] | ||
| assert next_draws == [saved.sample_from_clusters().tolist() for _ in range(3)] |
There was a problem hiding this comment.
The resume test passes an in-memory dictionary directly between samplers, so it cannot catch failures in the actual DeepSpeedEngine.save_checkpoint()/load_checkpoint() path, serialization, dataloader restoration, or distributed execution. Because this change modifies checkpoint-resume behavior observable by a training loop, add an integration test that compares uninterrupted and checkpoint-resumed sampling on actual devices and report the hardware result.
AGENTS.md reference: AGENTS.md:L35-L36
Useful? React with 👍 / 👎.
| sampler = _curriculum_data_sampler(tmp_path) | ||
| for _ in range(3): | ||
| sampler.sample_from_clusters() | ||
| assert sampler.state_dict()[CURRICULUM_LEARNING_NP_RNG_STATE] == sampler.np_rng.bit_generator.state |
There was a problem hiding this comment.
This assertion pins the test to the sampler's private np_rng.bit_generator.state representation rather than the checkpoint-resume contract; an otherwise correct implementation that copies, normalizes, or encodes the generator state differently would fail it. The following test already checks continued sampling behavior, so remove this assertion or replace it with a round-trip assertion over observable samples.
AGENTS.md reference: AGENTS.md:L30-L32
Useful? React with 👍 / 👎.
|
On the RNG assertion: It also pins something the resume test does not, which is which generator gets checkpointed. Putting the sampler back on the process-global On the engine-level resume test: |
|
#8425 opened a couple of hours ago with the same source change as this one, legacy-tuple branch included, so what actually differs is the tests. This PR adds three to Updated 2026-09-08: there are now three PRs on #8405, since #8460 opened this morning. All three make the same one-file change, saving and restoring Updated 2026-09-14: #8460 is closed, withdrawn by its own author, so the offer above no longer points anywhere. Two are left, this one and #8425, and only this one carries tests. The reason he gave was #8489, where the Data Efficiency library is named a candidate for deprecation for lack of bandwidth. If that is the direction, say so and I will close this. If it stays, the tests are the part worth keeping, since nothing else in the suite catches the sampler silently replaying its stream after a resume. |
1 participant
DeepSpeedDataSamplerdraws fromself.np_rng, the generator it seeds in__init__, butstate_dict()savednp.random.get_state()andload_state_dict()restored that. The global legacyRandomStatewas never the RNG the sampler used, so the generator stayed out of the checkpoint and a resumed run started its stream again from the seed: same cluster mix per step, same shuffles.Saves and restores
self.np_rng.bit_generator.stateinstead. A state dict written by an older version carries the old tuple, and that shape still goes throughnp.random.set_stateso old checkpoints keep loading the way they did.Added two tests for the resume path and one that loads an old-style state dict. The two behaviour ones fail on master.
Fixes #8405